package cn.doitedu;

import org.apache.hadoop.hive.ql.exec.UDAF;
import org.apache.hadoop.hive.ql.exec.UDAFEvaluator;
import org.roaringbitmap.longlong.Roaring64Bitmap;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.Serializable;
import java.nio.ByteBuffer;

public class AggBigIntToBitmap extends UDAF {


    public static class AggBigIntToBitmapEvaluator implements UDAFEvaluator {

        Acc acc;

        public AggBigIntToBitmapEvaluator() {
            System.out.println("构造了");
            init();
        }


        // 初始化中间累加器
        @Override
        public void init() {
            try {
                acc = new Acc();
                System.out.println("初始化了");
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }



        // 逐条聚合
        public boolean iterate(Long element) throws IOException {

            System.out.println("收到一条数据: " + element);
            Roaring64Bitmap bitmap = acc.getBitmap();
            bitmap.add(element);

            System.out.println("此时bitmap中的元素有： " + bitmap);

            acc.setBitmapBytes(SerDe.ser(bitmap));

            return true;
        }


        // map端输出的局部结果
        public Acc terminatePartial() {
            System.out.println("terminate partial ..........");
            return this.acc;
        }




        // reduce端合并map端输出的acc
        public boolean merge(Acc otherAcc) throws IOException {
            Roaring64Bitmap bm = this.acc.getBitmap();
            System.out.println("merge中，this: " + bm);
            System.out.println("merge中，other: " + otherAcc.getBitmap());
            bm.or(otherAcc.getBitmap());
            System.out.println("merge中，or: " + bm);

            this.acc.setBitmapBytes(SerDe.ser(bm));

            return true;
        }



        // 输出最终结果
        public byte[] terminate() throws IOException {
            System.out.println("terminate.......... = > " + this.acc.getBitmap());
            return this.acc.getBitmapBytes();
        }


    }



    public static class Acc implements Serializable {

        private byte[] bitmapBytes;

        public Acc() throws IOException {
            Roaring64Bitmap bm = Roaring64Bitmap.bitmapOf();
            this.bitmapBytes = SerDe.ser(bm);
        }

        public Roaring64Bitmap getBitmap() throws IOException {
            Roaring64Bitmap bm = SerDe.de(this.bitmapBytes);
            return bm;

        }

        public byte[] getBitmapBytes() {
            return bitmapBytes;
        }

        public void setBitmapBytes(byte[] bitmapBytes) {
            this.bitmapBytes = bitmapBytes;
        }

    }



    public static class SerDe{
        public static byte[] ser(Roaring64Bitmap bm) throws IOException {
            ByteArrayOutputStream bout = new ByteArrayOutputStream();
            DataOutputStream dout = new DataOutputStream(bout);
            bm.serialize(dout);

            return bout.toByteArray();
        }

        public static Roaring64Bitmap de(byte[] bytes) throws IOException {
            Roaring64Bitmap bm = Roaring64Bitmap.bitmapOf();
            bm.deserialize(ByteBuffer.wrap(bytes));
            return bm;
        }


    }

}
